import calendar
from datetime import datetime, timedelta, date
from typing import List
from decimal import Decimal


import numpy as np

from service.models import Month, Procedure, PredictionFact
from service.shedule_optimisation.coefficients import form_user_coefficients, form_events_coefficients
from service.shedule_optimisation.constraints import form_user_constraints, form_prediction_constraints, \
    form_events_constraints, form_week_work_events, count_event_hours
from service.shedule_optimisation.generator import generate_hours
from users.models import Specialization, User, SpecializationType
from service.models.event import Event
from tqdm import tqdm


def map_delta_hours_on_event(negative_events, hours_for_negative_events):
    return zip(negative_events, [sum(x) for x in hours_for_negative_events.T if sum(x) > 0])

def form_hours(user):
    if user.employment_fraction == Decimal('1.00'):
        return timedelta(hours=8)
    elif user.employment_fraction == Decimal('0.75'):
        return timedelta(hours=6)
    elif user.employment_fraction == Decimal('0.50'):
        return timedelta(hours=4)
    elif user.employment_fraction == Decimal('0.25'):
        return timedelta(hours=2)
    if user.employment_fraction == Decimal('0.10'):
        return timedelta(hours=2)


def get_week_days(month, week):
    res = calendar.monthrange(month.year, month.month)
    start_day = 1

    if date.fromisocalendar(month.year, week, 1).month != month.month:
        start_day = res[0] + 1

    end_day = 7

    if date.fromisocalendar(month.year, week, 7).month != month.month:
        end_day = date(month.year, month.month, res[1]).weekday() + 1

    if end_day > 5:
        end_day = 5

    return range(start_day, end_day + 1)


def get_delta_events(
        users: List[User],
        user_constraints,
        user_work_events,
        month,
        week,
) -> List[Event]:
    events = []
    for user in users:
        user_week_delta = 40 * user.employment_fraction - user_constraints[user.id]
        if user_week_delta > 0:
            for week_day in get_week_days(month, week):
                day = date.fromisocalendar(month.year, week, week_day).day
                a = user_work_events[user.id]
                day_events = list(filter(lambda x: x.day == day, a))
                if len(day_events) > 0:
                    day_event = day_events[0]
                    delta = 8 - count_event_hours(day_event)
                    if delta > 0:
                        events.append(
                            Event(
                                month=month,
                                kind='work_delta',
                                day=day,
                                user=user,
                                start_time=day_event.end_time,
                                end_time=(day_event.end_time + timedelta(hours=delta)).time()
                            )
                        )

                else:
                    start_time = datetime.strptime('09:00', '%H:%M')
                    d = form_hours(user)
                    end_time = start_time + d
                    events.append(
                        Event(
                            month=month,
                            kind='work_delta',
                            day=day,
                            user=user,
                            start_time=start_time.time(),
                            end_time=end_time.time()
                        )
                    )
    return events


def create_delta_events(negative_events_possibilities, month):
    delta_events = []
    for event_with_hours in negative_events_possibilities:
        if event_with_hours[1] > 0:
            event = event_with_hours[0]

            # Assuming event.start_time is a time object and event.date is a date object
            start_datetime = datetime.combine(event.date, event.start_time)
            end_datetime = start_datetime + timedelta(hours=event_with_hours[1])

            # If you need the time part of the end_datetime
            end_time = end_datetime.time()
            delta_events.append(
                Event(
                    month=month,
                    kind='work_delta',
                    day=event.day,
                    user=event.user,
                    start_time=event.start_time,
                    end_time=end_time
                )
            )
    Event.objects.bulk_create(delta_events)


class Optimiser:

    def __init__(self, month: Month):
        self.specialisations: List[Specialization] = list(Specialization.objects.all().order_by('specialization'))
        self.procedures: List[Procedure] = list(Procedure.objects.all())
        self.types:  List[SpecializationType] = list(SpecializationType.objects.all().select_related('specialization'))
        self.dict_types = {}
        self.month = month

        for x in self.specialisations:
            self.dict_types[x.specialization] = [y for y in self.types if y.specialization.specialization == x.specialization]

    def __get_events(self):
        return Event.objects.filter(month=self.month).all() \
            .select_related("month", "user", "user__main_specialization") \
            .prefetch_related("user__additional_specializations")

    def __get_predictions(self):
        return PredictionFact.objects.filter(month=self.month) \
            .select_related("specialization", "month")

    def __get_users(self):
        return User.objects.filter(groups__name='Врач').all() \
            .prefetch_related("additional_specializations") \
            .select_related("main_specialization")

    def optimise_work(self):
        events = self.__get_events()
        predictions = self.__get_predictions()
        users = self.__get_users()

        user_coefficients = form_user_coefficients(users, self.types, self.dict_types)

        for week in tqdm(set([prediction.week for prediction in predictions]), desc="Weeks"):
            week_predict = list(filter(lambda x: x.week == week, predictions))
            prediction_constraints = form_prediction_constraints(week_predict, self.procedures, self.types)

            user_work_events = form_week_work_events(users, events, week)
            user_constraints = form_user_constraints(users, user_work_events)

            hours_distribution = generate_hours(
                [coef[1] for coef in sorted(user_coefficients.items())],
                [hours[1] for hours in sorted(user_constraints.items())],
                prediction_constraints
            )

            spec_hours_distr = [sum(x) for x in hours_distribution]
            deltas = np.array(prediction_constraints) - np.array(spec_hours_distr)

            if sum(deltas) == 0:
                continue

            negative_events = get_delta_events(users, user_constraints, user_work_events, self.month, week)

            events_coefficients = form_events_coefficients(negative_events, user_coefficients)
            events_constraints = form_events_constraints(negative_events)

            hours_for_negative_events = generate_hours(
                events_coefficients,
                events_constraints,
                deltas
            )

            negative_events_possibilities = map_delta_hours_on_event(negative_events, hours_for_negative_events)
            create_delta_events(negative_events_possibilities, self.month)

            # hours_distr = np.array(hours_distribution)
            # doc_sum = np.sum(hours_distr, axis=0)
            # doc_distrib = hours_distr / doc_sum
            # spec_distribution = []
            # for user_index, col in enumerate(doc_distrib.T):
            #     for i, spec in enumerate(self.types):
            #         pf = list(filter(lambda x: x.specialization.id == spec.get('id'), predictions))[0]
            #         spec_distribution.append(
            #             SpecDistribution(week=pf, user=users[user_index], pct=round(col[i] * 100, 1)))
            #
            # SpecDistribution.objects.bulk_create(spec_distribution)
